package kafka

import (
	"context"
	"strings"
	"time"

	"github.com/Shopify/sarama"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

// ConsumeOffset 消费位置
var ConsumeOffset = sarama.OffsetNewest

// Handler 处理器
type Handler func(sarama.ConsumerGroupSession, *sarama.ConsumerMessage) error

// KafkaClient kafka客户端
type KafkaClient struct {
	client   sarama.Client
	producer sarama.SyncProducer
}

// NewKafkaClient 新建Kafka客户端
func NewKafkaClient(kafkaAddr string) (*KafkaClient, error) {
	c := sarama.NewConfig()

	c.Version = sarama.V2_1_0_0
	c.Producer.Return.Successes = true
	c.Producer.RequiredAcks = sarama.WaitForAll
	c.Producer.Partitioner = sarama.NewRoundRobinPartitioner
	c.Consumer.Offsets.Initial = ConsumeOffset

	if err := c.Validate(); err != nil {
		return nil, errors.WithStack(err)
	}

	client, err := sarama.NewClient(strings.Split(kafkaAddr, ","), c)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	producer, err := sarama.NewSyncProducerFromClient(client)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return &KafkaClient{
		client:   client,
		producer: producer,
	}, nil
}

// Client 获取客户端
func (k *KafkaClient) Client() sarama.Client {
	return k.client
}

// Consume 消费消息
func (k *KafkaClient) Consume(ctx context.Context, group, topic string, handler Handler) error {
	consumer, err := sarama.NewConsumerGroupFromClient(group, k.client)
	if err != nil {
		return errors.WithStack(err)
	}

	ch := &consumerHandler{handler: handler}

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			default:
				err := consumer.Consume(ctx, strings.Split(topic, ","), ch)
				if err == nil {
					continue
				}

				logrus.WithError(err).WithField("topic", topic).Error("consume message error")
				time.Sleep(time.Second)
			}
		}
	}()

	<-ctx.Done()
	err = consumer.Close()
	logrus.Debug("kafka consumer close")

	return errors.WithStack(err)
}

// SendMessage 发送消息
func (k *KafkaClient) SendMessage(message *sarama.ProducerMessage) error {
	_, _, err := k.producer.SendMessage(message)
	return errors.WithStack(err)
}

// Close 关闭
func (k *KafkaClient) Close() error {
	logrus.Debug("kafka close")
	if err := k.producer.Close(); err != nil {
		return errors.WithStack(err)
	}
	err := k.client.Close()
	return errors.WithStack(err)
}

// consumerHandler
type consumerHandler struct {
	handler Handler
}

// Setup setup
func (h *consumerHandler) Setup(session sarama.ConsumerGroupSession) error {
	return nil
}

// Cleanup cleanup
func (h *consumerHandler) Cleanup(session sarama.ConsumerGroupSession) error {
	return nil
}

// ConsumeClaim consume claim
func (h *consumerHandler) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
	for message := range claim.Messages() {
		if err := h.handler(session, message); err != nil {
			logrus.WithError(err).Error("handler kafka message error")
			return errors.WithStack(err)
		}
	}
	return nil
}
